AIP-72: Allow retrieving Variable from Task Context by amoghrajesh · Pull Request #45431 · apache/airflow

amoghrajesh · 2025-01-06T12:39:19Z

Summary of changes

Added a minimal Variable user-facing definition which will be used in DAG files by DAG authors
Added logic to get Variables in the context - both in "value" and "json" format

"value" is the raw form
"json" is the deserialised json form, we are trying to keep the contract between SDK and API server simple, they interact only in strings and the responsibility of serialising + deserialising lies on the client before sending it to the task sdk, not on the API server. This will enable multi language support too.

Object Glossary

-VariableResponse is auto-generated and tightly coupled with the API schema.
-VariableResult is runtime-specific and meant for internal communication between Supervisor & Task Runner.
-Variable class here is where the public-facing, user-relevant aspects are exposed, hiding internal details.

Testing

DAG:

from __future__ import annotations

from airflow.models.baseoperator import BaseOperator
from airflow.models.dag import dag


class CustomOperator(BaseOperator):
    def execute(self, context):
        import os
        os.environ["AIRFLOW_VAR_HI_MESSAGE"] = "hello_world"
        os.environ["AIRFLOW_VAR_JSON_VAR"] = "{\r\n  \"key1\": \"value1\",\r\n  \"key2\": \"value2\",\r\n  \"enabled\": true,\r\n  \"threshold\": 42\r\n}"
        task_id = context["task_instance"].task_id
        print(f"Hello World {task_id}!")
        print(context)
        print(context["var"]["value"].hi_message)
        print(context["var"]["json"].json_var)


@dag()
def var_from_context():
    CustomOperator(task_id="hello")


var_from_context()

This dag tests both the scenarios of a regular value context as well as json.

Case1: Variables found

Logs:

2c27ff9a5949
 ▶ Log message source details
2025-01-06 12:37:34.612546 [info     ] Filling up the DagBag from /files/dags/var_from_context.py [airflow.models.dagbag.DagBag]
2025-01-06 12:37:34.613168 [debug    ] Importing /files/dags/var_from_context.py [airflow.models.dagbag.DagBag]
2025-01-06 12:37:34.616934 [debug    ] Loaded DAG <DAG: var_from_context> [airflow.models.dagbag.DagBag]
2025-01-06 12:37:34.617131 [debug    ] DAG file parsed                [task] file=/files/dags/var_from_context.py
2025-01-06 12:37:34.641905 [warning  ] CustomOperator.execute cannot be called outside TaskInstance! [airflow.task.operators.unusual_prefix_c6632fd34e048ff55a9057c21ee5a54c16b99828_var_from_context.CustomOperator]
2025-01-06 12:37:34.642284 [info     ] Hello World hello!             [task] chan=stdout
2025-01-06 12:37:34.642393 [info     ] {'dag': <DAG: var_from_context>, 'inlets': [], 'map_index_template': None, 'outlets': [], 'run_id': 'manual__2025-01-06T12:37:34.037857+00:00', 'task': <Task(CustomOperator): hello>, 'task_instance': RuntimeTaskInstance(id=UUID('01943b9e-dadf-72eb-a47b-6c36e0fea3b7'), task_id='hello', dag_id='var_from_context', run_id='manual__2025-01-06T12:37:34.037857+00:00', try_number=1, map_index=-1, task=<Task(CustomOperator): hello>), 'ti': RuntimeTaskInstance(id=UUID('01943b9e-dadf-72eb-a47b-6c36e0fea3b7'), task_id='hello', dag_id='var_from_context', run_id='manual__2025-01-06T12:37:34.037857+00:00', try_number=1, map_index=-1, task=<Task(CustomOperator): hello>), 'var': {'json': <VariableAccessor (dynamic access)>, 'value': <VariableAccessor (dynamic access)>}, 'conn': <ConnectionAccessor (dynamic access)>, 'dag_run': DagRun(dag_id='var_from_context', run_id='manual__2025-01-06T12:37:34.037857+00:00', logical_date=datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), data_interval_start=datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), data_interval_end=datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), start_date=datetime.datetime(2025, 1, 6, 12, 37, 34, 492777, tzinfo=TzInfo(UTC)), end_date=None, run_type=<DagRunType.MANUAL: 'manual'>, conf={}), 'data_interval_end': datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), 'data_interval_start': datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), 'logical_date': datetime.datetime(2025, 1, 6, 12, 37, 34, 37857, tzinfo=TzInfo(UTC)), 'ds': '2025-01-06', 'ds_nodash': '20250106', 'task_instance_key_str': 'var_from_context__hello__20250106', 'ts': '2025-01-06T12:37:34.037857+00:00', 'ts_nodash': '20250106T123734', 'ts_nodash_with_tz': '20250106T123734.037857+0000'} [task] chan=stdout
2025-01-06 12:37:34.642268 [debug    ] Sending request                [task] json={"key":"hi_message","type":"GetVariable"}

2025-01-06 12:37:34.648998 [info     ] Variable(key='hi_message', value='hello_world', description=None) [task] chan=stdout
2025-01-06 12:37:34.649000 [debug    ] Sending request                [task] json={"key":"json_var","type":"GetVariable"}

2025-01-06 12:37:34.652506 [warning  ] Pydantic serializer warnings:
  Expected `str` but got `dict` with value `{'api_key': '12345', 'region': 'us-east-1'}` - serialized value may not be as expected [py.warnings] category=UserWarning filename=/usr/local/lib/python3.9/site-packages/pydantic/main.py lineno=426
2025-01-06 12:37:34.652586 [debug    ] Sending request                [task] json={"state":"success","end_date":"2025-01-06T12:37:34.652550Z","type":"TaskState"}

2025-01-06 12:37:34.652845 [info     ] Variable(key='json_var', value={'api_key': '12345', 'region': 'us-east-1'}, description=None) [task] chan=stdout

Case 2: Variable not found

Logs:

2c27ff9a5949
 ▶ Log message source details
2025-01-06 12:34:57.712062 [info     ] Filling up the DagBag from /files/dags/var_from_context.py [airflow.models.dagbag.DagBag]
2025-01-06 12:34:57.712579 [debug    ] Importing /files/dags/var_from_context.py [airflow.models.dagbag.DagBag]
2025-01-06 12:34:57.717120 [debug    ] Loaded DAG <DAG: var_from_context> [airflow.models.dagbag.DagBag]
2025-01-06 12:34:57.717434 [debug    ] DAG file parsed                [task] file=/files/dags/var_from_context.py
2025-01-06 12:34:57.743903 [warning  ] CustomOperator.execute cannot be called outside TaskInstance! [airflow.task.operators.unusual_prefix_c6632fd34e048ff55a9057c21ee5a54c16b99828_var_from_context.CustomOperator]
2025-01-06 12:34:57.744163 [info     ] Hello World hello!             [task] chan=stdout
2025-01-06 12:34:57.744270 [info     ] {'dag': <DAG: var_from_context>, 'inlets': [], 'map_index_template': None, 'outlets': [], 'run_id': 'manual__2025-01-06T12:34:56.856856+00:00', 'task': <Task(CustomOperator): hello>, 'task_instance': RuntimeTaskInstance(id=UUID('01943b9c-74e4-7b93-ba7e-7401dde98944'), task_id='hello', dag_id='var_from_context', run_id='manual__2025-01-06T12:34:56.856856+00:00', try_number=1, map_index=-1, task=<Task(CustomOperator): hello>), 'ti': RuntimeTaskInstance(id=UUID('01943b9c-74e4-7b93-ba7e-7401dde98944'), task_id='hello', dag_id='var_from_context', run_id='manual__2025-01-06T12:34:56.856856+00:00', try_number=1, map_index=-1, task=<Task(CustomOperator): hello>), 'var': {'json': <VariableAccessor (dynamic access)>, 'value': <VariableAccessor (dynamic access)>}, 'conn': <ConnectionAccessor (dynamic access)>, 'dag_run': DagRun(dag_id='var_from_context', run_id='manual__2025-01-06T12:34:56.856856+00:00', logical_date=datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), data_interval_start=datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), data_interval_end=datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), start_date=datetime.datetime(2025, 1, 6, 12, 34, 57, 589490, tzinfo=TzInfo(UTC)), end_date=None, run_type=<DagRunType.MANUAL: 'manual'>, conf={}), 'data_interval_end': datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), 'data_interval_start': datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), 'logical_date': datetime.datetime(2025, 1, 6, 12, 34, 56, 856856, tzinfo=TzInfo(UTC)), 'ds': '2025-01-06', 'ds_nodash': '20250106', 'task_instance_key_str': 'var_from_context__hello__20250106', 'ts': '2025-01-06T12:34:56.856856+00:00', 'ts_nodash': '20250106T123456', 'ts_nodash_with_tz': '20250106T123456.856856+0000'} [task] chan=stdout
2025-01-06 12:34:57.744177 [debug    ] Sending request                [task] json={"key":"hi_message","type":"GetVariable"}

2025-01-06 12:34:57.749149 [debug    ] Sending request                [task] json={"state":"failed","end_date":"2025-01-06T12:34:57.749108Z","type":"TaskState"}

TODO:

Writing variables from task SDK

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

task_sdk/src/airflow/sdk/execution_time/context.py

task_sdk/src/airflow/sdk/execution_time/supervisor.py

task_sdk/tests/api/test_client.py

kaxil

Few nits but lgtm

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>

amoghrajesh · 2025-01-06T19:14:35Z

Unrelated failure. Merging.

AIP-72: Allow retrieving Variable from Task Context

48a84ad

boring-cyborg bot added the area:task-sdk label Jan 6, 2025

amoghrajesh requested review from ashb and kaxil January 6, 2025 12:39

fixing a test

c6cf08b